from typing import List, Optional
import requests
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.embeddings import Embeddings
from core.config import settings
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class OllamaEmbeddings(Embeddings):
    def __init__(self, model: str = "llama2", base_url: str = "http://103.236.71.196:11434"):
        self.model = model
        self.base_url = base_url

    def _embed(self, text: str) -> List[float]:
        try:
            response = requests.post(
                f"{self.base_url}/api/embeddings",
                json={
                    "model": self.model,
                    "prompt": text
                },
                timeout=10  # 添加超时设置
            )
            response.raise_for_status()
            result = response.json()
            if "embedding" not in result:
                raise ValueError("Ollama response does not contain 'embedding' field")
            return result["embedding"]
        except Exception as e:
            logger.error(f"Ollama embedding error: {str(e)}")
            raise  # 重新抛出异常，让上层处理

    def embed_query(self, text: str) -> List[float]:
        return self._embed(text)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return [self._embed(text) for text in texts]


class AliEmbeddings(DashScopeEmbeddings):
    """阿里云百炼模型封装，便于统一调用"""

    def __init__(self, model: str = "tongyi-embedding-vision-flash", api_key: Optional[str] = None):
        super().__init__(
            model=model,
            dashscope_api_key=api_key or settings.ALI_EMBEDDING_API_KEY,
            max_retries=3
        )


class FallbackEmbeddings(Embeddings):
    """带有故障转移机制的嵌入模型，主模型失败时自动切换到备用模型"""

    def __init__(self, primary_embeddings: Embeddings, fallback_embeddings: Embeddings):
        self.primary_embeddings = primary_embeddings
        self.fallback_embeddings = fallback_embeddings

    def embed_query(self, text: str) -> List[float]:
        try:
            # 尝试调用主模型
            print("主模型调用开始")
            return self.primary_embeddings.embed_query(text)
        except Exception as e:
            print(f"主模型调用失败，切换到备用模型: {str(e)}")
            # 主模型失败，调用备用模型
            return self.fallback_embeddings.embed_query(text)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        try:
            # 尝试调用主模型
            print("主模型批量调用开始")
            return self.primary_embeddings.embed_documents(texts)
        except Exception as e:
            print(f"主模型批量调用失败，切换到备用模型: {str(e)}")
            # 主模型失败，调用备用模型
            return self.fallback_embeddings.embed_documents(texts)


def get_embeddings() -> Embeddings:
    """获取带有故障转移的嵌入模型实例"""
    # 初始化主模型（Ollama本地模型）
    primary_embeddings = OllamaEmbeddings(
        model=settings.EMBEDDING_MODEL_NAME,
        base_url=settings.EMBEDDING_URL
    )

    # 初始化备用模型（阿里云百炼）
    fallback_embeddings = AliEmbeddings(
        model="tongyi-embedding-vision-flash",

        api_key=settings.ALI_EMBEDDING_API_KEY
    )

    # 返回带有故障转移的嵌入模型
    return primary_embeddings
    # return FallbackEmbeddings(primary_embeddings, fallback_embeddings)
    # return fallback_embeddings

